# app/models/llm_agent.py
import json
from typing import List, Dict, Any, Union, Optional
from langchain.chat_models import ChatOpenAI
from langchain.agents import Tool, initialize_agent, AgentType
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.agents import AgentAction, AgentFinish
from langchain.schema import HumanMessage
from loguru import logger
from app.config import Config
from openai import OpenAI
class MaxIterationsReached(Exception):
    """自定义异常，用于终止代理执行"""
    pass

class CustomCallbackHandler(BaseCallbackHandler):
    def __init__(self):
        self.thoughts = []
        self.tool_usage_count = 0
        self.knowledge_base_used = False
        self.final_result = ""

    def on_agent_action(self, action: AgentAction, **kwargs):
        logger.info(f"######### 思考第{self.tool_usage_count}次")
        logger.info(f"######### 当前工具: {action.tool}")
        self.thoughts.append(f"Thought: {action.log}")
        self.tool_usage_count += 1

        if action.tool == "Knowledge Base":
            self.knowledge_base_used = True

    def on_agent_finish(self, finish: AgentFinish, **kwargs):
        # 处理返回值格式变化
        if isinstance(finish.return_values, dict):
            self.final_result = finish.return_values.get("output", "")
        else:
            self.final_result = str(finish.return_values)
        logger.info(f"######### 最终结果: {self.final_result}")

class EarlyStoppingCallback(BaseCallbackHandler):
    """通过回调控制最大轮次，替代max_iterations"""

    def __init__(self, max_turns: int):
        self.max_turns = max_turns
        self.current_turns = 0

    def on_agent_action(self, action: AgentAction, **kwargs) -> Any:
        self.current_turns += 1
        if self.current_turns >= self.max_turns:
            logger.info(f"######### 达到最大思考限制，停止思考 tools：{action.tool}")
            # 抛出异常终止代理执行
            raise MaxIterationsReached(f"达到最大思考轮次: {self.max_turns}")

class LLMAgent:
    def __init__(self, llm_config: Dict[str, Any], kb_tool: Tool, db_tool: Tool):
        self.llm_config = llm_config
        self.kb_tool = kb_tool
        self.db_tool = db_tool
        self.agent = None
        self.max_turns = Config.MAX_TURNS
        self.callback_handler = CustomCallbackHandler()
        self._initialize_agent()

    def _initialize_agent(self) -> None:
        logger.info(f"################ 初始化agent ################")
        llm = ChatOpenAI(**self.llm_config)
        tools = [self.kb_tool, self.db_tool]
        # 移除max_iterations，完全通过回调控制轮次
        self.agent = initialize_agent(
            tools=tools,
            llm=llm,
            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            handle_parsing_errors=True,
            verbose=False,
        )
        logger.info(f"################ agent初始化完毕 ################")

    def summarize_thoughts(self, query: str, thoughts: List[str]) -> str:
        """独立使用LLM生成总结，不通过代理"""
        logger.info(f"################ 思考完毕，整理结论 ################")
        prompt = f"""
        用户问题: {query}
        已进行{len(thoughts)}步思考，达到最大轮次限制。
        中间思考步骤: {" ".join(thoughts)}
        请基于以上内容生成最终回答，若信息不足请说明进展。
        """
        llm = ChatOpenAI(**self.llm_config)
        result = llm.invoke([HumanMessage(content=prompt)]).content  # 使用invoke方法
        logger.info(f"################ 结论整理完毕 ################")
        return result

    def judge_intention(self, query: str) -> bool:
        """判断问题意图是否为数据库操作，使用通用模型"""
        client = OpenAI(
            # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
            api_key="sk-079868c30081417fb54343c6919ba60d",
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )

        judge_prompt = f"用户问题: {query}。请判断这个问题是否需要进行数据库操作？请回答是或否。"
        completion = client.chat.completions.create(
            # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
            model="qwen-plus",
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": judge_prompt },
            ],
        )
        try:
            result = json.loads(completion.model_dump_json()).get("choices")[0].get("message").get("content")
        except:
            result = "否"
        return "是" in result

    def run(self, query: str) -> Dict[str, Any]:
        logger.info(f"################ 向agent提问：{query} ################")
        try:
            if not self.agent:
                self._initialize_agent()

            # 重置回调处理器
            self.callback_handler = CustomCallbackHandler()
            early_stopping = EarlyStoppingCallback(self.max_turns)

            # 判断问题意图
            is_db_operation = self.judge_intention(query)

            if '是' not in is_db_operation:
                # 非数据库操作，直接访问知识库结合模型能力生成答案
                kb_result = self.kb_tool.func(query)
                prompt = f"用户问题: {query}。知识库查询结果: {kb_result}。请基于以上内容生成最终回答。"
                llm = ChatOpenAI(**self.llm_config)
                result = llm.invoke([HumanMessage(content=prompt)]).content
            else:
                # 数据库操作，访问知识库获取内容，生成sql语句查询数据库，结合二者内容和模型能力生成答案
                kb_result = self.kb_tool.func(query)
                sql_prompt = f"用户问题: {query}。知识库查询结果: {kb_result}。请生成一个用于查询test库的SQL语句。只返回从SELECT到```的内容"
                llm = ChatOpenAI(**self.llm_config)
                sql = llm.invoke([HumanMessage(content=sql_prompt)]).content
                db_result = self.db_tool.func(sql)
                final_prompt = f"用户问题: {query}。知识库查询结果: {kb_result}。数据库查询结果: {db_result}。请基于以上内容生成最终回答。"
                result = llm.invoke([HumanMessage(content=final_prompt)]).content

            self.callback_handler.final_result = result

            response = {
                "thoughts": self.callback_handler.thoughts,
                "final_result": self.callback_handler.final_result,
                "tool_usage_count": self.callback_handler.tool_usage_count,
                "max_turns_reached": self.callback_handler.tool_usage_count >= self.max_turns,
                "knowledge_base_used": self.callback_handler.knowledge_base_used
            }
            logger.info(f"################ agent处理完成 ################")
            return response
        except Exception as e:
            logger.error(f"运行代理时出错: {e}")
            return {
                "error": str(e),
                "thoughts": [],
                "final_result": None,
                "tool_usage_count": 0,
                "max_turns_reached": False
            }